import torch
from agents.AgentMSAC import MSACAgent
from config.cfg import Config
from environment.ElectricScheduleEnv import env
from utils.misc import calc_output, split_action


def load_saved_model(args: Config):
    msac = MSACAgent(
        nagents=2,
        net_dims=[256, 256, 256],
        state_dim=env.state_dim,
        action_dim=env.action_dim,
        args=args,
    )
    actors_net_dict = torch.load(args.model_path)
    for a in msac.agents:
        a.load_params(actors_net_dict)
    return msac


if __name__ == '__main__':
    args = Config()
    if args.is_save_model:
        msac = load_saved_model(args)
        # 耗费的新能源出力
        cost_NE_A, cost_NE_B = 0, 0
        # 实际的新能源出力
        actual_NE_A, actual_NE_B = 0, 0
        state = env.reset()
        cntA, cntB = 0, 0
        for ts in range(env.max_step):
            actual_NE_A += env.A.Wind.ActualOutput[ts] + \
                env.A.Solar.ActualOutput[ts]
            actual_NE_B += env.B.Wind.ActualOutput[ts] + \
                env.B.Solar.ActualOutput[ts]
            print(f"========Hour[{ts}]========")
            actions = msac.take_target_actions(state)
            actionA, actionB = split_action(
                actions,
                actions_dims_dict={
                    'fire': [env.A.FireStationsSize, env.B.FireStationsSize],
                    "storage": [env.A.StorageStationsSize, env.B.StorageStationsSize],
                    "total": env.action_dim
                },
                ts=ts,
            )
            cost_NE_A += env.A.Wind.PredictOutput[ts] + \
                env.A.Solar.PredictOutput[ts]
            cost_NE_B += env.B.Wind.PredictOutput[ts] + \
                env.B.Solar.PredictOutput[ts]

            loadA, loadB = env.A.Load[ts], env.B.Load[ts]
            outputA, outputB = calc_output(
                action=actionA), calc_output(action=actionB)
            if outputA <= loadA:
                cntA += 1
            if outputB <= loadB:
                cntB += 1
            print(loadA, loadB)
            print(outputA, outputB)

            next_state, r, d = env.step(actionA, actionB)
            state = next_state
        print("Area A New energy consumption=%.2f\n" %
              (float(cost_NE_A)/float(actual_NE_A)))
        print("Area B New energy consumption=%.2f\n" %
              (float(cost_NE_B)/float(actual_NE_B)))
        print("Area A stability=%.2f\n" % (1-float(float(cntA)/float(24))))
        print("Area B stability=%.2f\n" % (1-float(float(cntB)/float(24))))
    else:
        print("model is not saved! Please change config to save model.")
